#https://zhuanlan.zhihu.com/p/17675153472
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import text
from sqlalchemy import create_engine
from mysqlconn import getSession
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import text
from sqlalchemy import create_engine
from mysqlconn import getSession
from sqlalchemy.ext.automap import automap_base
 
from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings 
hf = HuggingFaceEmbeddings(
        model_name='C:/Users/Administrator/.cache/modelscope/hub/models/Jerry0/text2vec-base-chinese' 
    )
 
#postgresql://postgres:123456@localhost:5432/hzai
dbHost = 'postgresql://postgres:123456@localhost:5432/hzai'
engine = create_engine(
    dbHost,
    echo=True,  # 是否打印SQL
    pool_size=10,  # 连接池的大小，指定同时在连接池中保持的数据库连接数，默认:5
    max_overflow=20,  # 超出连接池大小的连接数，超过这个数量的连接将被丢弃,默认: 5
)
Base = automap_base() 
Base.prepare(engine, reflect=True)

emb1 = Base.classes.emb1


Session = sessionmaker(bind=engine)
session = Session()
vectors = [
   '广东工业大学',
   "中山大学",
   "北京交通大学"
]
# 插入向量数据
for vector in vectors:
    query_result = hf.embed_query(vector)
    new_user = emb1(text='Alice',embedding=query_result)
    result=session.add(new_user) 
# Insert a new user

target_vector = hf.embed_query("广工")
# 构建查询语句
query = f"""
SELECT id, text,embedding, cosine_distance(embedding,  '{target_vector}') AS similarity
FROM emb1
ORDER BY similarity ASC
LIMIT 10;
"""
query=text(query)
print(query)
# 执行查询
result =session.execute(query )
# 打印结果
for row in result:
    print(f"ID: {row[0]}, text: {row[1]}, Similarity: {row[3]}")

 